package com.da.embedding;

import com.da.result.model.Embedding;
import com.da.tools.Request;
import com.da.tools.json.JsonObject;
import com.da.tools.json.JsonParser;

import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * 词向量生成
 *
 * @author da
 * @time 2024/7/12 下午 3:56
 */
public class OllamaEmbedding {
    /**
     * ollama 请求地址
     */
    private final String url;
    /**
     * 编码模型
     */
    private final String model;
    /**
     * json 解析器
     */
    private final JsonParser jsonParser = new JsonParser();

    /**
     * 构造函数
     *
     * @param url   ollama 请求地址
     * @param model 编码模型
     */
    public OllamaEmbedding(String url, String model) {
        this.url = url;
        this.model = model;
    }

    /**
     * 构造函数
     *
     * @param model 编码模型
     */
    public OllamaEmbedding(String model) {
        this("http://localhost:11434/api/embeddings", model);
    }

    /**
     * 构造函数
     */
    public OllamaEmbedding() {
        this("nomic-embed-text");
    }

    /**
     * 生成词向量
     *
     * @param prompt 输入文本
     * @return 词向量
     */
    public Embedding embedding(String prompt) {
        Request request = createRequest(prompt);
        return handleResponse(request);
    }

    /**
     * 处理请求响应
     *
     * @param request 请求对象
     * @return 词向量
     */
    @SuppressWarnings("unchecked")
    private Embedding handleResponse(Request request) {
        String response = request.send();
        if (request.getCode() == 200) {
            JsonObject data = (JsonObject) jsonParser.parse(response);
            return new Embedding((List<Double>) data.get("embedding"));
        }
        //保底返回一个空值
        return new Embedding();
    }

    /**
     * 创建请求对象
     *
     * @param prompt 输入文本
     * @return 请求对象
     */
    private Request createRequest(String prompt) {
        Map<String, Object> data = new HashMap<>();
        data.put("model", this.model);
        data.put("prompt", prompt);
        return new Request(this.url, data);
    }
}
